import asyncio
import json
import logging
import time
import websockets

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("ConsumerTest")

SAMPLE_RATE = 16000
VAD_FRAME_MS = 20
VAD_SILENCE_SEC = 6.0

# 扩展处理器，新增 REPORT_READY
handlers = {
    "QUESTION": lambda d: logger.info(f"📢 QUESTION: {d['text']}"),
    "TRANSCRIPT": lambda d: logger.info(
        f"{'💬[FINAL]' if d.get('is_final') else '✍[PARTIAL]'} {d.get('text')}"
    ),
    "DONE":    lambda d: logger.info("🏁 DONE"),
    "ERROR":   lambda d: logger.error(f"❌ ERROR: {d.get('message')}"),
    "INTERVIEW_COMPLETE": lambda d: logger.info("🎉 INTERVIEW_COMPLETE"),
    "REPORT_READY": lambda d: logger.info(f"📊 REPORT_READY: {json.dumps(d.get('report'), ensure_ascii=False)}"),
}

async def run_round(ws, pcm_file, first_round=False):
    if first_round:
        # 第一次主动拉题
        await ws.send(json.dumps({"event": "REQUEST_QUESTION"}))
        # 等 QUESTION
        while True:
            data = json.loads(await ws.recv())
            handlers.get(data["event"], lambda d: None)(data)
            if data["event"] == "QUESTION":
                break
    else:
        # 等后端自动下发的新题
        while True:
            data = json.loads(await ws.recv())
            handlers.get(data["event"], lambda d: None)(data)
            if data["event"] == "QUESTION":
                break

    # 启动 STT
    await ws.send(json.dumps({"event": "START_STT"}))
    # （可选）等 0.1–0.2s，或监听“STT_CONNECTED”反馈

    # 朗读完成
    await ws.send(json.dumps({"event": "QUESTION_READ_COMPLETE"}))

    # 发送音频
    chunk_size = int(SAMPLE_RATE * VAD_FRAME_MS/1000 * 2)
    frame_duration = VAD_FRAME_MS / 1000
    sent_bytes = 0
    t0 = time.time()
    with open(pcm_file, "rb") as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk: break
            await ws.send(chunk)
            sent_bytes += len(chunk)
            # 保持真实音频速率
            expected_time = (sent_bytes // chunk_size) * frame_duration
            await asyncio.sleep(max(0, expected_time - (time.time() - t0)))

    # 静音触发 DONE
    silent_frames = int((VAD_SILENCE_SEC + 1) / frame_duration)
    logger.info(f"⏳ Sending {silent_frames} silent frames")
    for _ in range(silent_frames):
        await ws.send(b"\x00"*chunk_size)
        await asyncio.sleep(frame_duration)

    # 等 DONE
    while True:
        data = json.loads(await ws.recv())
        handlers.get(data["event"], lambda d: None)(data)
        if data["event"] == "DONE":
            break


async def main(interview_id, pcm1, pcm2):
    uri = f"ws://localhost:8000/ws/interview/{interview_id}/"
    async with websockets.connect(uri, ping_interval=None) as ws:
        # … 两轮问答 …
        # 第一轮
        await run_round(ws, pcm1, first_round=True)
        # 第二轮
        await run_round(ws, pcm2, first_round=False)
        # 1) 结束面试
        await ws.send(json.dumps({"event": "END_INTERVIEW"}))

        # 2) 等 INTERVIEW_COMPLETE
        while True:
            data = json.loads(await ws.recv())
            handlers.get(data["event"], lambda d: None)(data)
            if data["event"] == "INTERVIEW_COMPLETE":
                logger.info("🕐 等待 REPORT_READY …")
                break

        # 3) 继续等 REPORT_READY，给后台任务留最多 30 秒
        start = time.time()
        while time.time() - start < 30:
            data = json.loads(await ws.recv())
            handlers.get(data["event"], lambda d: None)(data)
            if data["event"] == "REPORT_READY":
                logger.info("✅ 已收到 REPORT_READY，测试结束")
                return

        logger.error("⏰ 等待 REPORT_READY 超时")

if __name__ == "__main__":
    asyncio.run(main(1, "intor.pcm", "test1.pcm"))
